import random
import numpy as np
from mxnet import nd
from mxnet import gluon
from mxnet.gluon.data import DataLoader
from mxnet.gluon.utils import split_and_load
from .transform import random_color_distort
import utils

__all__ = ['MsDataLoader', 'split_and_load']

class FaceDetDataset(gluon.data.vision.ImageRecordDataset):
    """Face Detection dataset loaded from record file.
   
    Parameters
    ----------
    filename : str
        Path of the record file. It require both *.rec and *.idx file in the same
        directory, where raw image and labels are stored in *.rec file for better
        IO performance, *.idx file is used to provide random access to the binary file.

    Examples
    --------
    >>> face_dataset = FaceDetDataset('train.rec')
    >>> img, label = face_dataset[0]
    >>> print(img.shape, label.shape)
    (45, 45, 3) (1, 4)

    """
    def __init__(self, filename):
        super(FaceDetDataset, self).__init__(filename)
        self.anchors = None

    def __getitem__(self, idx):
        img, label = super(FaceDetDataset, self).__getitem__(idx)
        label = self._transform_label(label)
        img, label = self._data_augment(img, label)
        return img, utils.get_targets(self.anchors, label)
    
    def _transform_label(self, label):
        label = np.array(label).ravel()
        # format header(2 4), label(x1,y1,x2,y2,.....)
        assert int(label[0]) == 2
        assert int(label[1]) == 4
        label = label[2:].reshape(-1, 4)
        assert label.size > 0 # at least one face
        return label

    def _data_augment(self, img, label):
        img = random_color_distort(img)
        img, label = self._random_flip(img, label)
        return img, label

    def _random_flip(self, img, label):
        if random.random() < 0.5:
            img = nd.flip(img, axis=1)
            width = img.shape[1]
            label[:, (2,0)] = width - label[:, (0,2)]
            label[:, 1] = 1.0 - label[:, 3]
        return img, label

def MsDataLoader(net, rec_path, train, batch_size, num_workers):
    dataset = FaceDetDataset(rec_path)
    h, w = dataset[0][0].shape[:2]
    cls_maps, box_maps = net(nd.zeros((1,3,h,w)))
    dataset.anchors = utils.gen_anchors(cls_maps, [4,8], [16,24], [21,45])
    if train:
        return DataLoader(dataset, batch_size, shuffle=True, num_workers=num_workers, last_batch='rollover')
    else:
        return DataLoader(dataset, batch_size, num_workers=num_workers, last_batch='keep')
